import argparse
import os
import pickle

import tensorflow as tf
from tensorflow import keras

from constants import PROJECT_ROOT
import datetime

from minio import Minio
from minio.error import ResponseError
from time import gmtime, strftime


def train(data_dir: str):
    print(strftime("%z", gmtime()))

    # Training
    print("构建model")
    model = keras.Sequential([
        keras.layers.Flatten(input_shape=(28, 28)),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(10)])

    model.compile(optimizer='adam',
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

    print("开始加载训练集")
    with open(os.path.join(data_dir, 'train_images.pickle'), 'rb') as f:
        train_images = pickle.load(f)

    with open(os.path.join(data_dir, 'train_labels.pickle'), 'rb') as f:
        train_labels = pickle.load(f)

    print("开始训练")
    model.fit(train_images, train_labels, epochs=10)

    print("开始加载测试集")
    with open(os.path.join(data_dir, 'test_images.pickle'), 'rb') as f:
        test_images = pickle.load(f)

    with open(os.path.join(data_dir, 'test_labels.pickle'), 'rb') as f:
        test_labels = pickle.load(f)

    # Evaluation
    print("测试集效果评估")
    test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)

    print(f'Test Loss: {test_loss}')
    print(f'Test Acc: {test_acc}')

    # Save model

    model_path = os.path.join(PROJECT_ROOT, f'mnist-first.h5')
    print("保存模型:", model_path)
    tf.saved_model.save(model, model_path)
    print("开始上传模型到云端")

    time_stamp = datetime.datetime.now()
    print(time_stamp.strftime('%Y.%m.%d-%H:%M:%S'))

    minioClient = Minio('192.168.174.130:9000',
                        access_key='Q3AM3UQ867SPQQA43P2F',
                        secret_key='zuf+tfteSlswRu7BJ86wekitnifILbZam1KYY3TG',
                        secure=False)

    bucket_name = "model-save"
    try:
        if not minioClient.bucket_exists(bucket_name):
            minioClient.make_bucket(bucket_name)
        minioClient.fput_object(bucket_name, 'first-model', model_path)
    except ResponseError as err:
        print(err)

    buckets = minioClient.list_buckets()

    for bucket in buckets:
        print(bucket.name, bucket.creation_date)

    print(f'Model written to: {model_path}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Kubeflow FMNIST training script')
    parser.add_argument('--data_dir', help='path to images and labels.')
    args = parser.parse_args()

    train(data_dir=args.data_dir)
